/*
  Copyright (C) 2011, Kyungjoo Kim
  
  This file is part of LINAL (LINear ALgebra)
  
  All rights reserved.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.
  
  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.
  
  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
  Also add information on how to contact you by electronic and paper mail.
  
  
  Kyungjoo Kim
  iamkyungjoo@gmail.com
*/
#include "linal/common.hxx"
#include "linal/util.hxx"

#include "linal/format.hxx"

#include "linal/device/host/operation.hxx"

namespace linal      {
  namespace host   {
    namespace internal {
    int chol_nopiv(int uplo, FLA_Obj &A ) {
      FLA_Chol( uplo, A );
      return true;
    }
    int chol_piv_l(FLA_Obj &A, FLA_Obj &p) {
      LINAL_ERROR( A.m == A.n, LINAL_ERROR_A_SQUARE_MATRIX );

      FLA_Obj ATL,   ATR,      A00,  a01,      A02,
        ABL,   ABR,      a10,  alpha11,  a12,
        A20,  a21,      A22;

      FLA_Obj TTL,   TTR,      T00,  t01,      T02,
        TBL,   TBR,      t10,  tau11,    t12,
        T20,  t21,      T22;

      FLA_Obj pt,    pb,       p0,   rho1,     p2;

      FLA_Part_2x2( A, &ATL, &ATR,
                    &ABL, &ABR, 0, 0, FLA_TL );

      FLA_Part_2x1( p, &pt,
                    &pb, 0, FLA_TOP );

      while ( FLA_Obj_length( ATL ) < FLA_Obj_length( A ) ){
        FLA_Repart_2x2_to_3x3( ATL, ATR,     &A00, &a01,     &A02,
                               &a10, &alpha11, &a12,
                               ABL, ABR,     &A20, &a21,     &A22,
                               1, 1, FLA_BR );

        FLA_Repart_2x1_to_3x1( pt, &p0,
                               &rho1,
                               pb, &p2, 1, FLA_BOTTOM );

        //------------------------------------------------------------
        // ** Find pivot index
        FLA_Obj AL, AR;

        // ** Temporary vector on AR to store diagonal entrees
        FLA_Part_1x2( ABR, &AL, &AR, 1, FLA_RIGHT );
        get_diag( ABR, AR );

        // ** Find max index
        FLA_Amax( AR, rho1 );
        int *idx = (int*)(FLA_Obj_buffer_at_view( rho1 ));

        if (*idx) {

          // partition A again with target index
          FLA_Part_2x2( A, &TTL, &TTR,
                        &TBL, &TBR,
                        A00.m+(*idx), A00.n+(*idx), FLA_TL );

          FLA_Repart_2x2_to_3x3( TTL, TTR,     &T00, &t01,   &T02,
                                 &t10, &tau11, &t12,
                                 TBL, TBR,     &T20, &t21,   &T22,
                                 1, 1, FLA_BR );

          // ** Hermitize
          FLA_Copyt( FLA_CONJ_TRANSPOSE, a10, a01 );
          FLA_Copyt( FLA_CONJ_TRANSPOSE, a21, a12 );
          FLA_Copyt( FLA_CONJ_TRANSPOSE, t10, t01 );
          FLA_Copyt( FLA_CONJ_TRANSPOSE, t21, t12 );

          // ** diagonal pivot
          apply_pivots( LINAL_RIGHT, LINAL_TRANSPOSE,    rho1, ABR  );
          apply_pivots( LINAL_LEFT,  LINAL_NO_TRANSPOSE, rho1, ABR  );

          apply_pivots( LINAL_LEFT,  LINAL_NO_TRANSPOSE, rho1, ABL );
        }
        // ** Chol unblock algorithms
        LINAL_ERROR( FLA_SUCCESS == FLA_Sqrt( alpha11 ),
                     ">> Fail to do Sqrt");

        FLA_Inv_scal( alpha11, a21 );
        FLA_Her     ( FLA_LOWER_TRIANGULAR,
                      FLA_MINUS_ONE, a21, A22 );

        //------------------------------------------------------------
        FLA_Cont_with_3x3_to_2x2( &ATL,  &ATR,  A00, a01,     A02,
                                  a10, alpha11, a12,
                                  &ABL, &ABR,   A20, a21,     A22,
                                  FLA_TL );

        FLA_Cont_with_3x1_to_2x1( &pt,   p0,
                                  rho1,
                                  &pb,   p2, FLA_TOP );

      }

      return true;
    }

    int chol_piv_u(FLA_Obj &A, FLA_Obj &p) {
      LINAL_ERROR(false, LINAL_ERROR_NOT_YET_IMPLEMENTED);
      return true;
    }
    int chol_piv(int uplo, FLA_Obj &A, FLA_Obj &p) {

      switch ( uplo ) {
      case LINAL_UPPER_TRIANGULAR: chol_piv_u( A, p ); break;
      case LINAL_LOWER_TRIANGULAR: chol_piv_l( A, p ); break;
      default: LINAL_ERROR( false, LINAL_ERROR_CHECK_UPLO );       break;
      }
      return true;
    }
  }
}
}
